import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import random

def build_dataset(words):
    X, Y = [], []
    for w in words:
            #print(w)
            context = [0] * block_size
            for ch in w + '.':
                idx = stoi[ch]
                X.append(context)
                Y.append(idx)
                #print(''.join(itos[i] for i in context), '--->', itos[idx])
                context = context[1:] + [idx]

    X =   torch.tensor(X)
    Y =   torch.tensor(Y)
    print(X.shape, Y.shape)
    return X, Y

def draw_historgram(Layers):
    plt.figure(figsize=(20,4))
    legends = []
    for i, layer in enumerate(Layers):
        if isinstance(layer, Tanh):
            t = layer.out
            print("layer %d (%10s): mean %+.2f, std %.2f, saturated: %.2f" % (i, layer.__class__, t.mean(), t.std(), (t > 0.97).float().mean() * 100)) 
            hy, hx = torch.histogram(t, density=True)
            plt.plot(hx[:-1].tolist(), hy.tolist())
            legends.append(f'layer {i}')
    plt.legend(legends)
    plt.title("activation histogram")
    plt.show()

def draw_histogram_grad(Layers):
    plt.figure(figsize=(20,4))
    legends = []
    for i, layer in enumerate(Layers):
        if isinstance(layer, Tanh):
            t = layer.out.grad
            print("layer %d (%10s): mean %+.2f, std %.2f" % (i, layer.__class__, t.mean(), t.std()))
            hy, hx = torch.histogram(t, density=True)
            plt.plot(hx[:-1].tolist(), hy.tolist())
            legends.append(f'layer {i}')
    plt.legend(legends)
    plt.title("gradient histogram")
    plt.show()

def draw_2dim_histogram(parameters):
    plt.figure(figsize=(20,4))
    legends = []
    for i, p in enumerate(parameters):
        if p.ndim == 2:
            t = p.grad
            #print("weight(%10s): mean %+f, std %.e, grand:data ratio %e" % (tuple(p.shape), t.mean(), t.std(), t.std() / p.std()))
            hy, hx = torch.histogram(t, density=True)
            plt.plot(hx[:-1].tolist(), hy.tolist())
            legends.append(f'{i} {tuple(p.shape)}')
    plt.legend(legends)
    plt.title("gradient histogram")
    plt.show()

def draw_ud_histogram(ud, parameters):
    plt.figure(figsize=(20,4))
    legends = []
    for i, p in enumerate(parameters):
        plt.plot([ud[j][i] for j in range(len(ud))])
        legends.append('param %d' % i)
    plt.plot([0, len(ud)], [-3, -3], 'k') # these ratios shold be ~1e03
    plt.legend(legends)
    plt.title("gradient histogram")
    plt.show()

@torch.no_grad()
def split_loss(split):
    x,y = {
        'train': (X, Y),
        'val': (Xval, Yval),
        'test': (Xtest, Ytest)
    }[split]

    logits = model(x)
    loss = F.cross_entropy(logits, y)
    print(split, loss.item())




class Linear:
    def __init__(self, fan_in, fan_out, bias=True, generator=None):
        self.W = torch.randn(fan_in, fan_out, generator=generator) / fan_in**0.5
        self.b = torch.zeros(fan_out) if bias else None
    def __call__(self, x):
        self.out = x @ self.W
        if self.b is not None:
            self.out += self.b
        return self.out
    def parameters(self):
        return [self.W] if self.b is None else [self.W, self.b]
    
class BatchNorm1d:

    def __init__(self, dim, eps=1e-5, momentum=0.1):
        self.eps = eps
        self.momentum = momentum
        self.training = True
        #  parameters (trained with backprop)
        self.gamma = torch.ones(dim)
        self.beta = torch.zeros(dim)
        #  buffers (trained with a running 'momentum')
        self.running_mean = torch.zeros(dim)
        self.running_var = torch.ones(dim)
    def __call__(self, x):
        if x.ndim == 2:
            dim = 0
        elif  x.ndim == 3:
            dim = (0,1)
        if self.training:
            xmean = x.mean(dim, keepdim=True)
            xvar = x.var(dim, keepdim=True)
        else:
            xmean = self.running_mean
            xvar = self.running_var
        xhat = (x - xmean) / torch.sqrt(xvar + self.eps)
        self.out = self.gamma * xhat + self.beta
        if self.training:
            with torch.no_grad():
                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean
                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar

        return self.out
    def parameters(self):
        return [self.gamma, self.beta]
    
class Tanh:
    def __init__(self):
        self.out = None
    def __call__(self, x):
        self.out = torch.tanh(x)
        return self.out
    def parameters(self):
        return []

class Embedding:
    def __init__(self, num_embeddings, embedding_dim, generator=None):
        self.weight = torch.randn(num_embeddings, embedding_dim, generator=generator) / embedding_dim**0.5
        self.out = None
    def __call__(self, x):
        self.out = self.weight[x]
        return self.out

    def parameters(self):
        return [self.weight]

class FlattenConsecutive:
    def __init__(self, n):
        self.n = n

    def __call__(self, x):
        B, T, C  = x.shape
        x = x.view(B, T//self.n, C*self.n)
        if x.shape[1] == 1:
            x = x.squeeze(1)

        self.out = x
        return self.out
    def parameters(self):
        return []

class Sequential:
    def __init__(self, layers):
        self.layers = layers

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        self.out = x
        return self.out

    def parameters(self):
        params = []
        for layer in self.layers:
            ps = layer.parameters()
            params.extend(ps)
        return params

if __name__ == '__main__':
    words = open('names.txt', 'r').read().splitlines()
    print(len(words))
    chars = sorted(list(set(''.join(words))))
    chars = ['.'] + chars
    stoi = {s:i for i,s in enumerate(chars)}
    itos = {s:i for i,s in stoi.items()}    
    vocab_size = len(stoi)
    n_embed = 24  # dimension of the embedding
    n_hidden = 128 # number of neurons in hidden layers
    block_size = 8
    n = len(words)
    n1 = int (0.8*n)
    n2 = int (0.9*n)
    X, Y = build_dataset(words[:n1])
    Xval, Yval = build_dataset(words[n1:n2])
    Xtest, Ytest = build_dataset(words[n2:])
    g = torch.Generator().manual_seed(42)
    model = Sequential([
        Embedding(vocab_size, n_embed),
        FlattenConsecutive(2), Linear(n_embed * 2, n_hidden, bias=False, generator=g), BatchNorm1d(n_hidden), Tanh(),
        FlattenConsecutive(2), Linear(n_hidden * 2, n_hidden, bias=False, generator=g), BatchNorm1d(n_hidden), Tanh(),
        FlattenConsecutive(2), Linear(n_hidden * 2, n_hidden, bias=False, generator=g), BatchNorm1d(n_hidden), Tanh(),
        Linear(          n_hidden, vocab_size, generator=g)
    ])

    with torch.no_grad():
        model.layers[-1].W  *= 0.1
    parameters = model.parameters()

    print("number of parameters in total:",sum(p.nelement() for p in parameters)) #  number of parameters in total
    
    for p in parameters:
        p.requires_grad = True

    max_steps = 200000  
    batch_size = 32
    lossi = []
    ud = []
    
    for i in range(max_steps):
        #  minibatch construct
        idx = torch.randint(0, X.shape[0], (batch_size,)) 
        Xb = X[idx]
        Yb = Y[idx]

        #  forward pass
        logits = model(Xb)
        loss = F.cross_entropy(logits, Yb) #  loss function
        #for layer in model.layers:
        #    print(layer.__class__.__name__, ':', tuple(layer.out.shape))
        #   backward pass
        #for layer in Layers:
        #    layer.out.retain_grad() 
        
        for p in parameters:
            p.grad = None

        loss.backward()

        lr = 0.1 if i < 100000 else 0.01
        for p in parameters:
            p.data += -lr * p.grad

        if i % 10000 == 0:
            print(f'iter {i:7d}/{max_steps:7d}: {loss.item():.4f}')
        with torch.no_grad():
            ud.append([(lr*p.grad.std()/p.data.std()).log10().item() for p in parameters])
        lossi.append(loss.log10().item())
        
    plt.plot(lossi)
    plt.show()
    for layer in model.layers:
        layer.training = False
    
    split_loss('train')
    split_loss('val')
    #draw_historgram(Layers)
    #draw_histogram_grad(Layers)
    #draw_2dim_histogram(parameters)
    #draw_ud_histogram(ud, parameters)

